#!/usr/bin/env bash
set -euo pipefail

# Basic config
export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}
# Workaround: disable TorchDynamo to avoid FlashAttention varlen + Dynamo FakeTensor issues
export TORCHDYNAMO_DISABLE=${TORCHDYNAMO_DISABLE:-1}
# Optional: improve CUDA allocator stability on long runs
export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}

BASE_DIR=${BASE_DIR:-"/public/home/jianghaitenggroup/jianglab/WorkSpace/Qwen25_VL_FT"}

# MODEL_PATH=${MODEL_PATH:-"/public/home/jianghaitenggroup/jianglab/WorkSpace/Qwen2.5-VL-3B-Instruct/"}
MODEL_PATH=${MODEL_PATH:-"/public/home/jianghaitenggroup/jianglab/WorkSpace/Qwen25_VL_FT/saved_20250921/train/checkpoint-1400/"}
TRAIN_FILE=${TRAIN_FILE:-"${BASE_DIR}/saved/data/train.jsonl"}
VAL_FILE=${VAL_FILE:-"${BASE_DIR}/saved/data/val.jsonl"}
IMAGE_ROOT=${IMAGE_ROOT:-"/public/home/jianghaitenggroup/jianglab/WorkSpace/MASS-SS3_Wave_224_448_V2/"}
OUTPUT_DIR=${OUTPUT_DIR:-"${BASE_DIR}/saved/train/"}

PER_DEVICE_TRAIN_BS=${PER_DEVICE_TRAIN_BS:-1}
PER_DEVICE_EVAL_BS=${PER_DEVICE_EVAL_BS:-1}
GRAD_ACC=${GRAD_ACC:-16}
EPOCHS=${EPOCHS:-2}
LR=${LR:-1e-5}  # 2e-7 ~ 1e-6
WD=${WD:-0.1}
WARMUP=${WARMUP:-0.03}
EVAL_STEPS=${EVAL_STEPS:-200}
SAVE_STEPS=${SAVE_STEPS:-200}
SEED=${SEED:-42}
NUM_WORKERS=${NUM_WORKERS:-4}

# Pixels controls number of visual tokens
MIN_PIXELS=${MIN_PIXELS:-200704}   # 256*28*28
MAX_PIXELS=${MAX_PIXELS:-1003520}  # 1280*28*28

echo "Launching training on GPUs: $CUDA_VISIBLE_DEVICES)"

# Ensure versioned output dir exists
mkdir -p "${OUTPUT_DIR}"

accelerate launch --multi_gpu --num_processes 4 \
  ./train_val_loss.py \
  --model_path "$MODEL_PATH" \
  --train_file "$TRAIN_FILE" \
  --val_file "$VAL_FILE" \
  --image_root "$IMAGE_ROOT" \
  --output_dir "$OUTPUT_DIR" \
  --epochs $EPOCHS \
  --per_device_train_batch_size $PER_DEVICE_TRAIN_BS \
  --per_device_eval_batch_size $PER_DEVICE_EVAL_BS \
  --gradient_accumulation_steps $GRAD_ACC \
  --learning_rate $LR \
  --weight_decay $WD \
  --warmup_ratio $WARMUP \
  --eval_steps $EVAL_STEPS \
  --save_steps $SAVE_STEPS \
  --seed $SEED \
  --num_workers $NUM_WORKERS \
  --min_pixels $MIN_PIXELS \
  --max_pixels $MAX_PIXELS \
  --tensorboard
